# %%
# 自动求导
# 看一下训练的步骤吧
import torch
import torchvision

# 拿出预训练好的模型,resnet18
model = torchvision.models.resnet18(pretrained=True)
# 看一下模型的样子
print(model)
# 生成一些随机数据
data = torch.rand(1, 3, 64, 64)
labels = torch.rand(1, 1000)

# %%
# 这是进行前向传播
prediction = model(data)  # forward pass
# 看一下输出的样子
print(prediction.shape)
# %%
# 计算误差,也就是标签和输出的差距，注意loss都是求和的
# loss是一个标量，然后才好对参数进行求导
# 这个是直接误差，好像没有这个做法，，一般来说是MAE或者MSE
loss = (prediction - labels).sum()
# 然后反向传播,计算梯度
loss.backward()
# %%
# 看下参数的.grad就有东西了
param = next(iter(model.parameters()))
print(param.grad)

# %%
# 当然要使用优化器来更权重，加载SGD优化器，要传入模型参数（必须）
optim = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)

# %%
# 启动梯度下降
optim.step()
# %%
"""一般来说pytorch训练的通用流程"""
# 定义模型
# 配置优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
# 梯度清零
optimizer.zero_grad()
# 前向传播和反向传播
loss_fn(model(input), target).backward()
# 更新权重
optimizer.step()
# %%
"""微分"""
import torch

# 定义tensor，并且要进行求梯度
a = torch.tensor([2., 3], requires_grad=True)
b = torch.tensor([6., 4], requires_grad=True)
# 定义函数
Q = 3 * a ** 3 - b ** 2

# %%
# 根据上面，我们假设a,b是参数 Q是损失函数
# 我们需要做的是用Q对a,b求偏导
# 注意：误差需要被求和或者求平均后再进行求梯度
# 这样就是标量对参数求梯度，梯度会相对比较简单
Q.sum().backward()

# %%
# 打印出梯度
print(a.grad)
print(b.grad)
# %%
# 计算图可以将一些参数冻结
import torch

x = torch.rand(5, 5)
y = torch.rand(5, 5)
z = torch.rand((5, 5), requires_grad=True)
# 对于像rand这样的函数，会把第一个元组当成形状
# 所以上面的两种传参方式都是对的，加不加括号都一样
# rand(5,5) 是等价于 rand((5,5))的
# %%
# 定义计算，x,y都不要求梯度，z要求梯度
# 所有a不求梯度，而b要求梯度
a = x + y
b = x + z
print(a.requires_grad)
print(b.requires_grad)

# %%
# 冻结参数常用于调整预训练网络
# 来将resnet18的前面层次保留，将分类器重新训练
# 思路是先冻结全部参数，然后替换分类器层
from torch import nn, optim
# 加载模型
model = torchvision.models.resnet18(pretrained=True)
# 先冻结全部的参数
for param in model.parameters():
    param.requires_grad = False
# %%
# 再替换线性分类器层，新替换的参数默认不被冻结
# fc:full_connect
model.fc = nn.Linear(512, 10)

# %%
# 定义优化器,这时候只需要传入全连接层的参数就好了，只会更新全连接的参数
optimizer = optim.SGD(model.fc.parameters(), lr=1e-2, momentum=0.9)
# torch.no_grad()也可以实现同样的效果